"""
小红书视频搜索脚本
用于根据关键词搜索小红书视频并提取视频信息
"""

import asyncio
import sys
import os
import json
import re

# Add the project root directory to the Python path
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from playwright.async_api import async_playwright


class XiaohongshuVideoSearcher:
    """
    小红书视频搜索器，用于搜索和提取小红书视频信息
    """
    
    def __init__(self):
        """
        初始化小红书视频搜索器
        """
        self.playwright = None
        self.context = None
        self.page = None

    async def start_browser(self, headless=False):
        """
        启动浏览器
        """
        self.playwright = await async_playwright().start()
        self.context = await self.playwright.chromium.launch_persistent_context(
            user_data_dir="./user_data/xiaohongshu",
            channel="chrome",
            headless=headless,
            args=[
                "--no-sandbox",
                "--disable-dev-shm-usage",
                "--disable-blink-features=AutomationControlled",
                "--disable-extensions"
            ],
            viewport={"width": 1920, "height": 1080}
        )
        self.page = await self.context.new_page()
        return self.context, self.page

    async def search_videos(self, keyword, max_count=50):
        """
        搜索小红书视频并提取信息
        
        Args:
            keyword (str): 搜索关键词
            max_count (int): 最大获取视频数量
            
        Returns:
            list: 视频信息列表
        """
        # 导航到小红书搜索页面
        search_url = f"https://www.xiaohongshu.com/search_result?keyword={keyword}&source=web_search_result_notes"
        await self.page.goto(search_url)
        await self.page.wait_for_timeout(5000)  # 等待页面加载
        
        # 检查是否需要登录
        login_element = await self.page.query_selector("text=登录后查看搜索结果")
        if login_element:
            print("检测到需要登录，请在浏览器中手动登录...")
            print("注意：小红书支持多种登录方式：")
            print("1. 手机号登录")
            print("2. 小红书扫码登录")
            print("3. 微信扫码登录")
            input("登录完成后按回车键继续...")
        
        # 滚动页面以加载更多视频
        videos = []
        scroll_count = 0
        max_scroll_count = max_count // 10 + 5  # 根据需要的视频数量估算滚动次数
        
        while len(videos) < max_count and scroll_count < max_scroll_count:
            # 执行滚动
            await self.page.evaluate("window.scrollTo(0, document.body.scrollHeight)")
            await self.page.wait_for_timeout(3000)  # 等待内容加载
            scroll_count += 1
            
            # 提取当前页面的视频信息
            current_videos = await self.page.evaluate(r'''() => {
                const videos = [];
                // 尝试多种选择器查找视频元素
                let videoItems = Array.from(document.querySelectorAll('a[href*="/discovery/item/"]'));
                
                // 如果上面的选择器没有找到元素，尝试备用选择器
                if (videoItems.length === 0) {
                    const backupItems = document.querySelectorAll('a[href*="/item/"]');
                    videoItems = Array.from(backupItems).filter(item => item.href.includes('/item/'));
                }
                
                videoItems.forEach((item, index) => {
                    if (index >= 100) return; // 只处理前100个视频
                    
                    const url = item.href;
                    // 从URL中提取视频ID
                    const urlParts = url.split('/');
                    const videoId = urlParts[urlParts.length - 1];
                    
                    const textContent = item.textContent || '';
                    
                    // 提取用户名
                    let username = '未知用户';
                    const usernameElement = item.querySelector('.author-name') || 
                                          item.querySelector('.nickname') || 
                                          item.querySelector('.user-name');
                    if (usernameElement) {
                        username = usernameElement.textContent.trim();
                    } else {
                        // 尝试从文本内容中提取用户名
                        const usernameMatch = textContent.match(/@[^\\s]+/);
                        if (usernameMatch) {
                            username = usernameMatch[0];
                        }
                    }
                    
                    // 提取观看数/点赞数
                    let views = '观看数未知';
                    const viewsElement = item.querySelector('.like-count') || 
                                        item.querySelector('.collect-count') || 
                                        item.querySelector('.data-count');
                    if (viewsElement) {
                        views = viewsElement.textContent.trim();
                    } else {
                        // 尝试从文本内容中提取观看数
                        const viewsMatch = textContent.match(/\\d+[.\\d]*\\s*[万十万]*/);
                        if (viewsMatch) {
                            views = viewsMatch[0];
                        }
                    }
                    
                    // 提取标题
                    let title = '';
                    // 查找标题元素
                    const titleElement = item.querySelector('.note-title') || 
                                        item.querySelector('.title') || 
                                        item.querySelector('img[alt]');
                    if (titleElement && titleElement.textContent) {
                        title = titleElement.textContent.trim();
                    } else if (titleElement && titleElement.alt) {
                        title = titleElement.alt.trim();
                    } else if (textContent) {
                        // 从文本内容中提取标题
                        const titlePattern = /[\\u4e00-\\u9fa5]+[^\\n]*?[!?！？]*[\\u4e00-\\u9fa5]/;
                        const match = textContent.match(titlePattern);
                        if (match) {
                            title = match[0].trim();
                        }
                    }
                    
                    // 如果仍然没有标题，则使用部分文本内容
                    if (!title && textContent) {
                        title = textContent.substring(0, 50).trim();
                    }
                    
                    videos.push({
                        video_id: videoId,
                        url: url,
                        title: title,
                        username: username,
                        views: views,
                        full_title: textContent
                    });
                });
                
                return videos;
            }''')
            
            # 合并视频信息并去重
            for video in current_videos:
                # 检查是否已存在该视频
                if not any(v.get('video_id') == video.get('video_id') for v in videos):
                    videos.append(video)
            
            if len(videos) >= max_count:
                break
        
        # 限制返回的视频数量
        return videos[:max_count]
    
    async def close(self):
        """
        关闭浏览器
        """
        if self.context:
            await self.context.close()
        if self.playwright:
            await self.playwright.stop()


async def main():
    """
    主函数，处理命令行参数并执行搜索
    """
    if len(sys.argv) < 2:
        print("使用方法: python xiaohongshu_video_search.py <关键词> [数量]")
        print("示例: python xiaohongshu_video_search.py 美妆教程 50")
        return
    
    keyword = sys.argv[1]
    max_count = 50  # 默认获取50个视频
    
    if len(sys.argv) > 2:
        try:
            max_count = int(sys.argv[2])
        except ValueError:
            print("数量参数必须是整数，使用默认值50")
    
    print(f"正在搜索关键词: {keyword}, 需要获取 {max_count} 个视频")
    
    # 创建搜索器实例
    searcher = XiaohongshuVideoSearcher()
    
    try:
        # 启动浏览器
        await searcher.start_browser(headless=False)
        
        # 搜索视频
        videos = await searcher.search_videos(keyword, max_count)
        
        # 输出JSON格式结果
        print(json.dumps(videos, ensure_ascii=False, indent=2))
            
    except Exception as e:
        print(f"搜索过程中发生错误: {e}")
    finally:
        # 关闭浏览器
        await searcher.close()


if __name__ == "__main__":
    asyncio.run(main())